package com.newx.auth.license.tool;

import cn.hutool.core.util.StrUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.security.*;
import java.security.spec.EncodedKeySpec;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Base64;
import java.util.zip.ZipEntry;
import java.util.zip.ZipOutputStream;

/**
 * 生成公钥和私钥
 * 生成 RSA 密钥对：使用 KeyPairGenerator 生成 2048 位的 RSA 密钥对。
 * 编码私钥为 PKCS#8 格式：将私钥编码为 PKCS#8 格式并保存到 .key 文件中。
 * 提取公钥并编码为 X.509 格式：从私钥中提取公钥并编码为 X.509 格式，保存到 .key 或 .pem 文件中。
 * PEM 文件格式：确保输出文件符合 PEM 文件格式（带有 -----BEGIN ... ----- 和 -----END ... ----- 包围的 Base64 编码字符串）。
 * @Created by yxl
 * @Date 2025/2/14 16:42
 */
public class KeyPairGeneratorTool {

    private static final Logger log = LoggerFactory.getLogger(KeyPairGeneratorTool.class);
    private static final String PRIVATE_KEY_FILENAME = "private.key";
    private static final String PUBLIC_KEY_FILENAME = "public.pem";

    /**
     * 生成公钥和私钥
     * @param filePath 文件保存路径
     * @return 成功消息
     * @throws IOException 如果文件操作失败
     * @throws NoSuchAlgorithmException 如果找不到算法
     */
    public static String generatorOpenSSL(String filePath) throws IOException, NoSuchAlgorithmException {
        if (StrUtil.hasBlank(filePath)) {
            throw new IllegalArgumentException("License file path cannot be blank.");
        }

        Path path = Paths.get(filePath).toAbsolutePath().normalize();
        Files.createDirectories(path);

        String normalizedPath = path.toString() + File.separator;

        // 生成密钥对
        KeyPair keyPair = generateKeyPair();
        PrivateKey privateKey = keyPair.getPrivate();
        PublicKey publicKey = keyPair.getPublic();

        // 编码并保存私钥到 .key 文件
        saveKey(privateKey, normalizedPath + PRIVATE_KEY_FILENAME, "PRIVATE");

        // 编码并保存公钥到 .pem 文件
        saveKey(publicKey, normalizedPath + PUBLIC_KEY_FILENAME, "PUBLIC");

        return "密钥对已成功生成并保存在路径：" + normalizedPath;
    }

    /**
     * 生成公钥和私钥，并返回打包为 ZIP 文件的字节数组
     * @return ZIP 文件的字节数组
     * @throws IOException 如果文件操作失败
     * @throws NoSuchAlgorithmException 如果找不到算法
     */
    public static byte[] generateAndPackKeysToZip() throws IOException, NoSuchAlgorithmException {
        // 生成密钥对
        KeyPair keyPair = generateKeyPair();
        PrivateKey privateKey = keyPair.getPrivate();
        PublicKey publicKey = keyPair.getPublic();

        // 在内存中创建 ZIP 文件
        try (ByteArrayOutputStream baos = new ByteArrayOutputStream();
             ZipOutputStream zos = new ZipOutputStream(baos)) {

            // 将私钥添加到 ZIP 文件
            addKeyToZip(zos, privateKey, PRIVATE_KEY_FILENAME, "PRIVATE");

            // 将公钥添加到 ZIP 文件
            addKeyToZip(zos, publicKey, PUBLIC_KEY_FILENAME, "PUBLIC");

            return baos.toByteArray();
        }
    }

    private static void addKeyToZip(ZipOutputStream zos, Key key, String entryName, String keyType) throws IOException {
        EncodedKeySpec encodedKeySpec = (key instanceof PrivateKey)
                ? new PKCS8EncodedKeySpec(key.getEncoded())
                : new X509EncodedKeySpec(key.getEncoded());

        byte[] encodedKey = Base64.getMimeEncoder().encode(encodedKeySpec.getEncoded());
        String encodedKeyString = chunkString(new String(encodedKey, StandardCharsets.US_ASCII), 64);

        ZipEntry zipEntry = new ZipEntry(entryName);
        zos.putNextEntry(zipEntry);
        zos.write(("-----BEGIN " + keyType + " KEY-----\n").getBytes(StandardCharsets.US_ASCII));
        zos.write(encodedKeyString.getBytes(StandardCharsets.US_ASCII));
        zos.write(("\n-----END " + keyType + " KEY-----\n").getBytes(StandardCharsets.US_ASCII));
        zos.closeEntry();
    }

    /**
     * 生成密钥对
     * @return 密钥对
     * @throws NoSuchAlgorithmException 如果找不到算法
     */
    private static KeyPair generateKeyPair() throws NoSuchAlgorithmException {
        KeyPairGenerator keyGen = KeyPairGenerator.getInstance("RSA");
        keyGen.initialize(2048); // 使用 2048 位密钥长度
        return keyGen.generateKeyPair();
    }

    /**
     * 保存密钥
     * @param key 密钥
     * @param fileName 文件名
     * @param keyType 密钥类型
     * @throws IOException 如果文件操作失败
     */
    private static void saveKey(Key key, String fileName, String keyType) throws IOException {
        EncodedKeySpec encodedKeySpec = (key instanceof PrivateKey)
                ? new PKCS8EncodedKeySpec(key.getEncoded())
                : new X509EncodedKeySpec(key.getEncoded());

        byte[] encodedKey = Base64.getMimeEncoder().encode(encodedKeySpec.getEncoded());
        String encodedKeyString = chunkString(new String(encodedKey, StandardCharsets.US_ASCII), 64);

        try (FileOutputStream fos = new FileOutputStream(fileName)) {
            fos.write(("-----BEGIN " + keyType + " KEY-----\n").getBytes(StandardCharsets.US_ASCII));
            fos.write(encodedKeyString.getBytes(StandardCharsets.US_ASCII));
            fos.write(("\n-----END " + keyType + " KEY-----\n").getBytes(StandardCharsets.US_ASCII));
        }
    }

    /**
     * 将字符串分割为指定长度的块
     * @param str 要分割的字符串
     * @param chunkSize 每个块的长度
     * @return 分割后的字符串
     */
    private static String chunkString(String str, int chunkSize) {
        StringBuilder result = new StringBuilder();
        for (int i = 0; i < str.length(); i += chunkSize) {
            if (i > 0) {
                result.append('\n');
            }
            int end = Math.min(str.length(), i + chunkSize);
            result.append(str.substring(i, end));
        }
        return result.toString();
    }
}

